package sortAlgorithm;

import java.util.ArrayList;
import java.util.List;

import com.alibaba.fastjson2.JSON;

/**
 * @description: <a href="https://leetcode.cn/problems/smallest-k-lcci/description/">面试题 17.14. 最小K个数</a> 使用快排的思路解决 topk 问题
 * @author: zhaoyinglin
 * @date: 2024/3/13
 */
public class SmallestKICCI {

    public static void main(String[] args) {

        int[] array = new int[] { 1, 3, 5, 7, 2, 4, 6, 8 };

        List<Integer> smallK = quickSort(array, 4);

        System.out.println(JSON.toJSONString(smallK));

        System.out.println(JSON.toJSONString(array));
    }

    private static List<Integer> quickSort(int[] arr, int k) {

        List<Integer> smallK = new ArrayList<>();

        helper(arr, 0, arr.length - 1, smallK, k);

        return smallK;
    }

    private static void helper(int[] arr, int start, int end, List<Integer> smallK, int k) {

        if (start > end) {
            return;
        }

        if (start == end) {
            if (start <= k - 1) {
                smallK.add(arr[start]);
            }
            return;
        }

        int p = partition(arr, start, end);

        if (p <= k - 1) {
            smallK.add(arr[p]);
        }

        helper(arr, start, p - 1, smallK, k);

        helper(arr, p + 1, end, smallK, k);
    }

    private static int partition(int[] array, int start, int end) {
        int p = end, i = start, j = start;

        for (; i < end; i++) {
            if (array[i] <= array[p]) {
                swap(i, j, array);
                j++;
            }
        }

        swap(p, j, array);

        return j;
    }

    private static void swap(int i, int j, int[] array) {
        int tmp = array[i];
        array[i] = array[j];
        array[j] = tmp;
    }

    class Solution {
        public int[] smallestK(int[] arr, int k) {

            if (arr == null || k >= arr.length) {
                return arr;
            }

            List<Integer> smallK = quickSort(arr, k);

            int[] smallKArray = new int[smallK.size()];

            for (int i = 0; i < smallK.size(); i++) {
                smallKArray[i] = smallK.get(i);
            }

            return smallKArray;
        }

        private List<Integer> quickSort(int[] arr, int k) {

            List<Integer> smallK = new ArrayList<>();

            helper(arr, 0, arr.length - 1, smallK, k);

            return smallK;
        }

        private void helper(int[] arr, int start, int end, List<Integer> smallK, int k) {

            if (start > end) {
                return;
            }

            if (start == end) {
                if (start <= k - 1) {
                    smallK.add(arr[start]);
                }
                return;
            }

            int p = partition(arr, start, end);

            if (p <= k - 1) {
                smallK.add(arr[p]);
            }

            helper(arr, start, p - 1, smallK, k);

            helper(arr, p + 1, end, smallK, k);
        }

        private int partition(int[] array, int start, int end) {
            int p = end, i = start, j = start;

            for (; i < end; i++) {
                if (array[i] <= array[p]) {
                    swap(i, j, array);
                    j++;
                }
            }

            swap(p, j, array);

            return j;
        }

        private void swap(int i, int j, int[] array) {
            int tmp = array[i];
            array[i] = array[j];
            array[j] = tmp;
        }

    }

}
